import os
import numpy as np
import pandas as pd
from scipy.io import loadmat
from typing import Dict, Any, Tuple
from ecgdetectors import Detectors
from scipy.signal import butter, filtfilt, medfilt
from pykalman import KalmanFilter  # 导入 pykalman 库
from scipy.stats import iqr

# Butterworth低通滤波器
def butter_lowpass_filter(data: np.ndarray, cutoff: float, fs: float, order: int) -> np.ndarray:
    nyquist = 0.5 * fs
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    y = filtfilt(b, a, data)
    return y

# Butterworth高通滤波器
def butter_highpass_filter(data: np.ndarray, cutoff: float, fs: float, order: int) -> np.ndarray:
    nyquist = 0.5 * fs
    normal_cutoff = cutoff / nyquist
    b, a = butter(order, normal_cutoff, btype='highpass', analog=False)
    y = filtfilt(b, a, data)
    return y

# 归一化信号
def normalization(data: np.ndarray) -> np.ndarray:
    """
    归一化信号，确保信号基线为0且幅度在[-1, 1]范围内。

    Parameters:
    - data: 输入信号

    Returns:
    - 归一化后的信号，基线为0，幅度范围为[-1, 1]
    """
    data = data - np.mean(data)  # 去除直流分量
    abs_max = np.max(np.abs(data))  # 获取最大绝对值
    if abs_max != 0:
        data = data / abs_max  # 归一化到[-1, 1]
    return data

def butter_bandpass_filter(data: np.ndarray, lowcut: float, highcut: float, fs: float, order: int) -> np.ndarray:
    """带通滤波器"""
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return filtfilt(b, a, data)

def preprocess_signal(signal: np.ndarray, fs: float) -> np.ndarray:
    """
    信号预处理：带通滤波 + 去除基线漂移
    """
    
    # 带通滤波（0.5-40Hz for ECG）
    signal = butter_bandpass_filter(signal, lowcut=0.5, highcut=40.0, fs=fs, order=2)
    
    # 去除基线漂移
    signal = remove_baseline_drift(signal, fs)
    
    return signal

def remove_baseline_drift(
    data: np.ndarray,
    fs: float,
    cutoff: float = 0.5,
    order: int = 2
) -> np.ndarray:
    """
    使用高通滤波器去除基线漂移。

    参数:
    - data: 输入的ECG信号
    - fs: 采样率
    - cutoff: 截止频率，默认0.5Hz
    - order: 滤波器阶数，默认2阶

    返回:
    - 去除基线漂移后的信号
    """
    return butter_highpass_filter(data, cutoff, fs, order)

def adaptive_kalman_filter_bank(ecg_signal: np.ndarray, fs: float) -> np.ndarray:
    """优化的自适应卡��曼滤波器组"""
    # 预处理
    ecg_signal = preprocess_signal(ecg_signal, fs)
    
    # 检测R峰和分离信号分量
    detectors = Detectors(fs)
    r_peaks = detectors.pan_tompkins_detector(ecg_signal)
    high_freq_signal, low_freq_signal = separate_ecg_components(ecg_signal, r_peaks, fs)
    
    # 处理高频分量
    high_freq_denoised = process_signal_component(high_freq_signal, fs)
    
    # 处理低频分量
    low_freq_denoised = process_signal_component(low_freq_signal, fs)
    
    # 合并结果
    denoised_signal = combine_ecg_components(high_freq_denoised, low_freq_denoised, r_peaks, fs)
    return normalization(denoised_signal)

def process_signal_component(signal: np.ndarray, fs: float) -> np.ndarray:
    """使用pykalman处理信号分量"""
    # 自适应估计初始参数
    signal_var, noise_var = adaptive_noise_estimation(signal)
    
    # 设置最小方差阈值
    min_variance = 1e-5
    median_signal_var = np.median(signal_var)
    median_noise_var = np.median(noise_var)
    
    # 确保方差不低于最小阈值
    median_signal_var = max(median_signal_var, min_variance)
    median_noise_var = max(median_noise_var, min_variance)
    
    # 配置卡尔曼滤波器
    kf = KalmanFilter(
        n_dim_state=1,
        n_dim_obs=1,
        initial_state_mean=0,
        initial_state_covariance=median_signal_var,
        transition_matrices=np.array([[1]]),
        observation_matrices=np.array([[1]]),
        transition_covariance=median_signal_var * 0.1,
        observation_covariance=median_noise_var
    )
    
    # EM算法估计参数
    # 原代码尝试解包返回值，导致错误
    # parameters, kf = estimate_parameters_em_with_pykalman(signal, num_iterations=5)
    kf = estimate_parameters_em_with_pykalman(signal, num_iterations=5)  # 仅获取 KalmanFilter 对象
    
    # 确保信号中不包含NaN值
    if np.isnan(signal).any():
        signal = np.nan_to_num(signal)
    
    # 应用滤波和平滑
    try:
        smoothed_state_means, smoothed_state_cov = kf.smooth(signal.reshape(-1, 1))
    except ValueError as e:
        # 处理滤波过程中可能出现的错误
        print(f"Kalman filter error: {e}")
        return signal  # 返回原始信号或进行其他处理
    
    return smoothed_state_means.flatten()

def em_parameter_estimation(kf: KalmanFilter, signal: np.ndarray, n_iter: int = 5) -> KalmanFilter:
    """优化的EM参数估计过程"""
    try:
        kf.em(signal, n_iter=n_iter)
    except:
        # 如果EM估计失败，保持原始参数
        pass
    return kf

def separate_ecg_components(ecg_signal: np.ndarray, r_peaks: np.ndarray, fs: float) -> Tuple[np.ndarray, np.ndarray]:
    """
    分离ECG信号为高频和低频成分
    """
    high_freq_signal = extract_high_freq(ecg_signal, r_peaks, fs)
    low_freq_signal = ecg_signal - high_freq_signal
    return high_freq_signal, low_freq_signal

def extract_high_freq(ecg_signal: np.ndarray, r_peaks: np.ndarray, fs: float) -> np.ndarray:
    """
    提取高频分量（QRS复合波）- 向量化优化
    """
    high_freq_signal = np.zeros_like(ecg_signal)
    qrs_width = int(0.1 * fs)  # QRS宽度约为100ms
    half_width = qrs_width // 2
    for peak in r_peaks:
        start = max(peak - half_width, 0)
        end = min(peak + half_width, len(ecg_signal))
        high_freq_signal[start:end] = ecg_signal[start:end]
    return high_freq_signal

def combine_ecg_components(denoised_high: np.ndarray, denoised_low: np.ndarray, r_peaks: np.ndarray, fs: float) -> np.ndarray:
    """
    合并高频和低频滤波后的信号 - 向量化优化
    """
    denoised_ecg = denoised_low.copy()
    qrs_width = int(0.1 * fs)  # QRS宽度约为100ms
    half_width = qrs_width // 2
    for peak in r_peaks:
        start = max(peak - half_width, 0)
        end = min(peak + half_width, len(denoised_ecg))
        denoised_ecg[start:end] = denoised_high[start:end]
    return denoised_ecg

def adaptive_noise_estimation(signal: np.ndarray, window_size: int = 100) -> Tuple[np.ndarray, np.ndarray]:
    """自适应估计信号和噪声方差（使用滑动窗口）"""
    half_window = window_size // 2

    # 使用Pandas的rolling窗口计算滑动中位数
    median_signal = pd.Series(signal).rolling(window=window_size, center=True, min_periods=1).median().values
    mad = np.median(np.abs(signal - median_signal)) * 1.4826
    signal_var = np.full(signal.shape, mad**2)

    # 使用Pandas的rolling窗口计算滑动IQR
    rolling_series = pd.Series(signal)
    q75 = rolling_series.rolling(window=window_size, center=True, min_periods=1).quantile(0.75).values
    q25 = rolling_series.rolling(window=window_size, center=True, min_periods=1).quantile(0.25).values
    iqr_values = q75 - q25
    noise_var = (iqr_values ** 2) / 1.349

    return signal_var, noise_var

def estimate_parameters_em_with_pykalman(
    signal: np.ndarray, 
    A_init: float = 1.0, 
    C_init: float = 1.0, 
    Q_init: float = 1e-5, 
    R_init: float = 1e-4, 
    num_iterations: int = 10
) -> KalmanFilter:
    """
    使用pykalman的EM方法对参数进行估计，并返回KalmanFilter对象。
    """
    # 将信号变形为 (N,1) 的观测数据
    observations = signal.reshape(-1, 1)

    # 初始化KalmanFilter
    kf = KalmanFilter(
        transition_matrices=np.array([[A_init]]),
        observation_matrices=np.array([[C_init]]),
        transition_covariance=np.array([[Q_init]]),
        observation_covariance=np.array([[R_init]]),
        initial_state_mean=np.array([0.0]),
        initial_state_covariance=np.array([[1.0]])
    )

    # EM迭代估计参数
    kf = kf.em(observations, n_iter=num_iterations)

    return kf  # 仅返回 kf

# ...existing code...

# 移除多余的kalman_filter_adaptive函数
# def kalman_filter_adaptive(
#     signal: np.ndarray,
#     A: float,
#     C: float,
#     Q: np.ndarray,
#     R: np.ndarray
# ) -> np.ndarray:
#     ...
    
# 移除多余的remove_baseline_drift函数，如果不再使用
# def remove_baseline_drift(
#     data: np.ndarray,
#     fs: float,
#     cutoff: float = 0.5,
#     order: int = 2
# ) -> np.ndarray:
#     ...